#!/usr/bin/env python3
import Angle
from pymodbus.client import ModbusSerialClient as ModbusClient
from pymodbus.register_write_message import WriteSingleRegisterRequest
import time

class Robot:
    def __init__(self, port='/dev/ttyS7', baudrate=115200, timeout=0.1):
        # Initialize Modbus RTU client
        self.client = ModbusClient(
            method='rtu',
            port=port,
            baudrate=baudrate,
            timeout=timeout
        )
        self.client.connect()
        
        self.motor_commands = {
            # 电机使能指令
            'enable': {
                'motor1': bytearray([0x01,0x06,0x00,0x3e,0x00,0x0f,0xa8,0x02]),
                'motor2': bytearray([0x06,0x06,0x00,0x3e,0x00,0x0f,0xa9,0xb5]),
                'joint1': bytearray([0x02,0x06,0x00,0x3e,0x00,0x0f,0xa8,0x31]),
                'joint2': bytearray([0x03,0x06,0x00,0x3e,0x00,0x0f,0xa9,0xe0]),
                'joint3': bytearray([0x04,0x06,0x00,0x3e,0x00,0x0f,0xa8,0x57]),
                'joint4': bytearray([0x05,0x06,0x00,0x3e,0x00,0x0f,0xa9,0x86])
            },
            # 电机禁用指令
            'disable': {
                'joint1': bytearray([0x02,0x06,0x00,0x3e,0x00,0x06,0x68,0x37]),
                'joint2': bytearray([0x03,0x06,0x00,0x3e,0x00,0x06,0x69,0xe6]),
                'joint3': bytearray([0x04,0x06,0x00,0x3e,0x00,0x06,0x68,0x51]),
                'joint4': bytearray([0x05,0x06,0x00,0x3e,0x00,0x06,0x69,0x80])
                # 其他电机禁用指令...
            },
            #模式切换指令  位置模式
            'position_mode':{
                'joint2': bytearray([0x02,0x06,0x00,0x3c,0x00,0x01,0x88,0x35]),
                'joint3': bytearray([0x03,0x06,0x00,0x3c,0x00,0x01,0x89,0xe4]),
                'joint4': bytearray([0x04,0x06,0x00,0x3c,0x00,0x01,0x88,0x53]),
                'joint5': bytearray([0x05,0x06,0x00,0x3c,0x00,0x01,0x89,0x82])
            },
            #速度模式
            'velocity_mode':{
                'joint2': bytearray([0x02,0x06,0x00,0x3c,0x00,0x03,0x09,0xf4]),
                'joint3': bytearray([0x03,0x06,0x00,0x3c,0x00,0x03,0x08,0x25]),
                'joint4': bytearray([0x04,0x06,0x00,0x3c,0x00,0x03,0x09,0x92]),
                'joint5': bytearray([0x05,0x06,0x00,0x3c,0x00,0x03,0x08,0x43])
            }
        }

    def send_command(self, unit, address, value):
        result = self.client.write_register(
            address=address,
            value=value,
            unit=unit
        )
        time.sleep(0.0025)
        return result

    def get_angle_data(self, num):
        command = Angle.SendOrder(num)
        response = self.client.execute(command)
        if not response or not response.isError():
            return Angle.OrderToAngle(response)
        return None

    def Position_mode(self):
        for motor in self.motor_commands['position_mode']:
            cmd = self.motor_commands['position_mode'][motor]
            self.send_command(cmd['unit'], cmd['address'], cmd['value'])

    def Velocity_mode(self):
        for motor in self.motor_commands['velocity_mode']:
            cmd = self.motor_commands['velocity_mode'][motor]
            self.send_command(cmd['unit'], cmd['address'], cmd['value'])

    def enable_motor(self, motor_name):
        if motor_name in self.motor_commands['enable']:
            cmd = self.motor_commands['enable'][motor_name]
            self.send_command(cmd['unit'], cmd['address'], cmd['value'])
        else:
            print(f"未知电机: {motor_name}")
    
    def disable_motor(self, motor_name):
        if motor_name in self.motor_commands['disable']:
            cmd = self.motor_commands['disable'][motor_name]
            self.send_command(cmd['unit'], cmd['address'], cmd['value'])
        else:
            print(f"未知电机: {motor_name}")
    
    def enable_all(self):
        for motor in self.motor_commands['enable']:
            self.enable_motor(motor)
    
    def disable_all(self):
        for motor in self.motor_commands['disable']:
            self.disable_motor(motor)
    
    def move_front(self, motor_name, speed):
        self.send_command(*Angle.SpeedToOrder(1, -speed))
        self.send_command(*Angle.SpeedToOrder(6, speed))
    
    def stop_move(self):
        self.send_command(*Angle.SpeedToOrder(1, 0))
        self.send_command(*Angle.SpeedToOrder(6, 0))

    def move_joint_angle(self, joint_name, angle):
        print(f"{joint_name}号电机转动到: {angle}")
        self.send_command(*Angle.AngleToOrder(joint_name, angle))

    def move_straight(self):
        for i in range(4):
            self.send_command(*Angle.AngleToOrder(i+2, 1))
            print(f"joint{i+2} moving to Zero")

    def move_straight1(self):
        for i in range(4):
            if i == 2:
                unit, address, value = Angle.AngleToOrder(i+2, -1)
            else:
                unit, address, value = Angle.AngleToOrder(i+2, 1)
            
            response = self.send_command(unit, address, value)
            
            if response:
                print(f"Response from joint {i+2}: {response}")
            else:
                print(f"No response from joint {i+2}")

    def move_joint_speed(self, joint_name, direction):
        speed = 1500 if direction == 1 else (-1500 if direction == -1 else 0)
        self.send_command(*Angle.SpeedToOrder(joint_name, speed))

    def close(self):
        self.client.close()
    
    def __del__(self):
        if hasattr(self, 'client'):
            self.client.close()